{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "e6870d75",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import statsmodels.api as sm\n",
    "from sklearn.metrics import mean_absolute_error, mean_squared_error\n",
    "import numpy as np\n",
    "from sklearn.preprocessing import StandardScaler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "390d28d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. 加载数据\n",
    "data = pd.read_csv('../data_process/big_data.csv')\n",
    "data['date'] = pd.to_datetime(data['date'])\n",
    "data.set_index('date', inplace=True)\n",
    "\n",
    "scaler = StandardScaler()\n",
    "scaler = scaler.fit(data)\n",
    "df_scaled = scaler.transform(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "cc8c5dae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((751, 12), (192, 12))"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 3. 分割训练集和测试集\n",
    "train = df_scaled[:-192]  # 使用数据集中除最后30天的部分作为训练集\n",
    "test = df_scaled[-192:]  # 使用数据集中最后30天的部分作为测试集\n",
    "\n",
    "train = pd.DataFrame(train, columns=['close','total_cases','new_cases_smoothed','total_deaths','new_deaths_smoothed','stringency_index','open','highest','lowest','neg','neu','pos'])\n",
    "test = pd.DataFrame(test, columns=['close','total_cases','new_cases_smoothed','total_deaths','new_deaths_smoothed','stringency_index','open','highest','lowest','neg','neu','pos'])\n",
    "\n",
    "train.shape, test.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "685f7f2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4. 定义外部解释变量（训练集和测试集）\n",
    "exog_train = train[['total_cases','new_cases_smoothed','total_deaths','new_deaths_smoothed','stringency_index','open','highest','lowest','neg','neu','pos']].dropna()\n",
    "exog_test = test[['total_cases','new_cases_smoothed','total_deaths','new_deaths_smoothed','stringency_index','open','highest','lowest','neg','neu','pos']].dropna()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "1312a206",
   "metadata": {},
   "outputs": [],
   "source": [
    "endog_train = train.loc[exog_train.index, 'close']\n",
    "endog_test = test.loc[exog_test.index, 'close']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "c52c7516",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/fanyc/opt/anaconda3/envs/tensorflow/lib/python3.8/site-packages/statsmodels/base/model.py:604: ConvergenceWarning: Maximum Likelihood optimization failed to converge. Check mle_retvals\n",
      "  warnings.warn(\"Maximum Likelihood optimization failed to \"\n"
     ]
    }
   ],
   "source": [
    "model = sm.tsa.ARIMA(endog=endog_train, exog=exog_train, order=(6, 1, 6))\n",
    "fit_model = model.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "098ae8ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "forecast = fit_model.forecast(steps=len(endog_test), exog=exog_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "10e286f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_forecast = fit_model.predict(start=exog_train.index[0], end=exog_train.index[-1], exog=exog_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0117e173",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_mae = mean_absolute_error(endog_train, train_forecast)\n",
    "train_mse = mean_squared_error(endog_train, train_forecast)\n",
    "train_rmse = np.sqrt(train_mse)\n",
    "train_mape = np.mean(np.abs((endog_train - train_forecast) / endog_train)) * 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "e5835401",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.9883367278779541, 0.002096044994785396, 0.03763090557877693)"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test_mae = mean_absolute_error(endog_test, forecasted_values)\n",
    "# test_mse = mean_squared_error(endog_test, forecasted_values)\n",
    "# test_rmse = np.sqrt(test_mse)\n",
    "# test_mape = np.mean(np.abs((endog_test - forecasted_values) / endog_test)) * 100\n",
    "from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_error\n",
    "r2_score(endog_test, forecast), mean_squared_error(endog_test, forecast), mean_absolute_error(endog_test, forecast)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "f67f60cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(751    0.996444\n",
       " 752    1.079637\n",
       " 753    1.107723\n",
       " 754    1.113704\n",
       " 755    1.129877\n",
       "          ...   \n",
       " 938   -0.389738\n",
       " 939   -0.418544\n",
       " 940   -0.318750\n",
       " 941   -0.488314\n",
       " 942   -0.480551\n",
       " Name: predicted_mean, Length: 192, dtype: float64,\n",
       " 0      0.973722\n",
       " 1      1.085783\n",
       " 2      1.125930\n",
       " 3      1.144899\n",
       " 4      1.163868\n",
       "          ...   \n",
       " 187   -0.430897\n",
       " 188   -0.444484\n",
       " 189   -0.318696\n",
       " 190   -0.456440\n",
       " 191   -0.552600\n",
       " Name: close, Length: 192, dtype: float64)"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "forecast,endog_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a93d8564",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "13.125637729669782\n"
     ]
    }
   ],
   "source": [
    "print(train_mae)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bcba9b4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}